#include "SrpcProvider.h"
#include "SrpcApplication.h"
#include "SrpcHeader.pb.h"
#include "SrpcLogger.h"
#include "ZooKeeperUtil.h"
#include <cstdlib>
#include <string>
#include <zookeeper/proto.h>

void SrpcProvider::RegistService(google::protobuf::Service *service) {
  ServiceInfo service_info;
  // 获取服务描述对象
  const google::protobuf::ServiceDescriptor *dsp = service->GetDescriptor();
  // 获取服务名称
  std::string service_name = dsp->name();
  // 获取服务端对象service的方法数量
  int method_count = dsp->method_count();
  SrpcLogger::INFO("service name: " + service_name +
                   " method count: " + std::to_string(method_count));

  // 遍历中的所有方法，并注册到服务信息中
  for (int i = 0; i < method_count; ++i) {
    // 获取服务中方法的描述
    const google::protobuf::MethodDescriptor *mdsp = dsp->method(i);
    std::string method_name = mdsp->name();
    SrpcLogger::INFO("method name: " + method_name);
    // 同一个服务下的所有方法存储在一个unordered_map中
    service_info.method_map.emplace(method_name, mdsp);
  }
  // 服务名称
  service_info.service = service;
  service_map.emplace(service_name, service_info);
}

void SrpcProvider::Run() {
  std::string ip =
      SrpcApplication::GetInstance().GetConfig().Load("rpcserverip");
  int port = std::atoi(
      SrpcApplication::GetInstance().GetConfig().Load("rpcserverport").c_str());
  // 使用muduo网路库，创建地址对象
  muduo::net::InetAddress address(ip, port);
  // 创建TcpServer对象
  std::shared_ptr<muduo::net::TcpServer> server =
      std::make_shared<muduo::net::TcpServer>(&event_loop, address,
                                              "SrpcProvider");
  // 绑定连接回调和消息回调
  server->setConnectionCallback(
      std::bind(&SrpcProvider::OnConnection, this, std::placeholders::_1));
  server->setMessageCallback(
      std::bind(&SrpcProvider::OnMessage, this, std::placeholders::_1,
                std::placeholders::_2, std::placeholders::_3));
  // 设置muduo库的线程数量
  server->setThreadNum(4);
  // 将服务注册到ZooKeeper上
  ZooKeeperClient zkCkient;
  zkCkient.Start();
  for (auto &sp : service_map) {
    std::string service_path = "/" + sp.first;
    // 创建服务节点
    zkCkient.Create(service_path.c_str(), "", 0);
    for (auto &mp : sp.second.method_map) {
      std::string method_path = service_path + "/" + mp.first;
      std::string method_path_data;
      method_path_data = ip + ":" + std::to_string(port);
      zkCkient.Create(method_path, method_path_data, ZOO_EPHEMERAL);
    }
  }
  SrpcLogger::INFO("RpcProvider start service at ip: " + ip +
                   " port: " + std::to_string(port));

  // 启动网络服务
  server->start();
  event_loop.loop();
}

void SrpcProvider::OnConnection(const muduo::net::TcpConnectionPtr &conn) {
  if (!conn->connected()) {
    conn->shutdown();
  }
}

void SrpcProvider::OnMessage(const muduo::net::TcpConnectionPtr &conn,
                             muduo::net::Buffer *buffer,
                             muduo::Timestamp receive_time) {
  // 读取请求内容
  std::string recv_buf = buffer->retrieveAllAsString();
  // 反序列化
  google::protobuf::io::ArrayInputStream raw_input(recv_buf.data(),
                                                   recv_buf.size());
  google::protobuf::io::CodedInputStream coded_input(&raw_input);
  uint32_t header_size{0};
  // 解析header_size
  coded_input.ReadVarint32(&header_size);
  // 根据header_size 读取数据头的原始字符流
  std::string rpc_header_str;
  Srpc::RpcHeader SrpcHeader;
  std::string service_name;
  std::string method_name;
  uint32_t args_size{0};
  // 设置读取限制
  google::protobuf::io::CodedInputStream::Limit msg_limit =
      coded_input.PushLimit(header_size);
  coded_input.ReadString(&rpc_header_str, header_size);
  // 恢复之前的限制，以便安全地继续读取其他数据
  coded_input.PopLimit(msg_limit);

  if (SrpcHeader.ParseFromString(rpc_header_str)) {
    service_name = SrpcHeader.service_name();
    method_name = SrpcHeader.method_name();
    args_size = SrpcHeader.args_size();
  } else {
    SrpcLogger::ERROR("SrpcHeader parse error");
    return;
  }
  // rpc参数
  std::string args_str;
  if (!coded_input.ReadString(&args_str, args_size)) {
    SrpcLogger::ERROR("read args error");
    return;
  }
  // 获取service对象和method对象
  auto it = service_map.find(service_name);
  if (it == service_map.end()) {
    SrpcLogger::ERROR(service_name + "is not exist!");
    return;
  }
  auto mit = it->second.method_map.find(method_name);
  if (mit == it->second.method_map.end()) {
    SrpcLogger::ERROR(service_name + "." + method_name + " is not exist!");
    return;
  }
  // 获取服务对象
  google::protobuf::Service *service = it->second.service;
  // 获取方法对象
  const google::protobuf::MethodDescriptor *method = mit->second;
  // 生成RPC方法调用请求的request和响应response
  // 动态创建请求对象
  google::protobuf::Message *request =
      service->GetRequestPrototype(method).New();
  if (!request->ParseFromString(args_str)) {
    SrpcLogger::ERROR(service_name + "." + method_name + " is parse error!");
    return;
  }
  // 动态创建响应对象
  google::protobuf::Message *response =
      service->GetResponsePrototype(method).New();
  // 绑定回调函数，用于在方法调用完成后发送响应
  google::protobuf::Closure *done =
      google::protobuf::NewCallback<SrpcProvider,
                                    const muduo::net::TcpConnectionPtr &,
                                    google::protobuf::Message *>(
          this, &SrpcProvider::SendRpcResponse, conn, response);
  // 调用对应方法
  service->CallMethod(method, nullptr, request, response, done);
}

void SrpcProvider::SendRpcResponse(const muduo::net::TcpConnectionPtr &conn,
                                   google::protobuf::Message *response) {
  std::string response_str;
  if (response->SerializeToString(&response_str)) {
    conn->send(response_str);
  } else {
    SrpcLogger::ERROR("serialize error!");
  }
  conn->shutdown();
}

SrpcProvider::~SrpcProvider() {
  SrpcLogger::INFO("~SrpcProvider()");
  event_loop.quit();
}